# 导入必要的库
import torch
import torch.nn.functional as F
import torch.optim as optim
from torch_geometric.data import Data
from torch_geometric.nn import GATConv
import networkx as nx
import matplotlib.pyplot as plt


####  使用图注意力网络（GAT）

# 数据预处理
bom_structure = {
    "显示屏": {
        "屏幕": ["玻璃", "元器件"],
        "主控制板": ["元器件", "玻璃"],
        "液晶模组": ["液晶面板", "背光模组"]
    },
    "键盘": {
        "键帽": [],
        "电路板": ["元器件", "金属丝"],
        "主控制板": ["元器件"],
    },
    "中央处理器": {
        "运算器": ["元器件"],
        "控制单元": ["主控制板", "元器件"],
        "寄存器": ["元器件"],
        "主控制板": ["元器件"]
    }
}

# 提取所有零件组件并建立索引映射
components = set(part for product, contents in bom_structure.items() for part in contents.keys())
components.update(
    subpart for product, contents in bom_structure.items() for subparts in contents.values() for subpart in subparts)
components = sorted(components)
component_to_index = {comp: index for index, comp in enumerate(components)}

# 构建边索引和节点特征
edge_list = [(component_to_index[part], component_to_index[subpart]) for product, contents in bom_structure.items() for
             part, subparts in contents.items() for subpart in subparts]
edge_index = torch.tensor(edge_list, dtype=torch.long).t().contiguous()
node_features = torch.eye(len(components))

# 创建 PyTorch Geometric 数据对象
graph_data = Data(x=node_features, edge_index=edge_index)

# 模型定义
class GNNModel(torch.nn.Module):
    def __init__(self, num_node_features, hidden_size, out_size):
        super(GNNModel, self).__init__()
        self.conv1 = GATConv(num_node_features, hidden_size)
        self.conv2 = GATConv(hidden_size, out_size)

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)
        return x

# 模型训练
hidden_size = 32
out_size = len(components)
model = GNNModel(node_features.size(1), hidden_size, out_size)

optimizer = optim.Adam(model.parameters(), lr=0.005)
model.train()
for epoch in range(500):
    optimizer.zero_grad()
    out = model(graph_data.x, graph_data.edge_index)
    loss = F.mse_loss(out, torch.rand((len(components), out_size)))
    loss.backward()
    optimizer.step()
    if epoch % 50 == 0:
        print(f"Epoch {epoch}: Loss {loss.item()}")

# 模型应用
def recommend_subcomponents(model, graph_data, parent_component_name):
    model.eval()
    with torch.no_grad():
        embeddings = model(graph_data.x, graph_data.edge_index)
        parent_index = component_to_index[parent_component_name]
        child_indices = graph_data.edge_index[1][graph_data.edge_index[0] == parent_index]
        parent_embedding = embeddings[parent_index].unsqueeze(0)
        child_embeddings = embeddings[child_indices]
        similarities = F.cosine_similarity(parent_embedding, child_embeddings)
        sorted_indices = similarities.argsort(descending=True)
        unique_indices = torch.unique(child_indices[sorted_indices])
        recommended_subcomponents = [components[i] for i in unique_indices.tolist()]
        recommended_scores = [similarities[sorted_indices[i]].item() for i in range(len(unique_indices))]
        return list(zip(recommended_subcomponents, recommended_scores))

# 示例
parent_component = "主控制板"
recommendations = recommend_subcomponents(model, graph_data, parent_component)
print(f"为{parent_component}推荐的子部件及相似度得分: {recommendations}")